package com.zcloud.eventstudydemo.util;

import com.hankcs.hanlp.HanLP;
import com.hankcs.hanlp.seg.common.Term;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 基于HanLP的关键词提取工具类
 */
public class HanLPKeywordExtractor {
    
    // 停用词列表
    private static final Set<String> STOP_WORDS = new HashSet<>(Arrays.asList(
        "的", "了", "在", "是", "我", "有", "和", "就", "不", "人", "都", "一", "一个", "上", "也", "很", "到", "说", "要", "去", "你", "会", "着", "没有", "看", "好", "自己", "这"
    ));
    
    /**
     * 使用HanLP从文本中提取关键词
     * @param text 输入文本
     * @param topK 返回关键词数量
     * @return 关键词列表
     */
    public static List<String> extractKeywords(String text, int topK) {
        if (text == null || text.isEmpty()) {
            return new ArrayList<>();
        }
        
        // 使用HanLP进行分词
        List<Term> termList = HanLP.segment(text);
        
        // 过滤停用词和不需要的词性
        List<String> words = termList.stream()
                .filter(term -> !STOP_WORDS.contains(term.word))
                .filter(term -> term.word.length() > 1)
                .filter(term -> term.nature.startsWith("n") || term.nature.startsWith("v") || 
                               term.nature.startsWith("a") || term.nature.startsWith("eng"))
                .map(term -> term.word)
                .collect(Collectors.toList());
        
        // 统计词频
        Map<String, Integer> wordFreq = new HashMap<>();
        for (String word : words) {
            wordFreq.put(word, wordFreq.getOrDefault(word, 0) + 1);
        }
        
        // 按词频排序并返回topK个关键词
        return wordFreq.entrySet().stream()
                .sorted(Map.Entry.<String, Integer>comparingByValue().reversed())
                .limit(topK)
                .map(Map.Entry::getKey)
                .collect(Collectors.toList());
    }
    
    /**
     * 使用HanLP进行关键词提取（基于TextRank算法）
     * @param text 输入文本
     * @param topK 返回关键词数量
     * @return 关键词列表
     */
    public static List<String> extractKeywordsByTextRank(String text, int topK) {
        if (text == null || text.isEmpty()) {
            return new ArrayList<>();
        }
        
        // 使用HanLP的TextRank算法提取关键词
        return HanLP.extractKeyword(text, topK);
    }
}